from ultralytics import YOLO
import airsim
import cv2
import numpy as np
import time

# 连接 AirSim 客户端
client = airsim.MultirotorClient()
client.confirmConnection()
print("Connected to AirSim!")

# 加载 YOLO 模型
model = YOLO("YOLO/yolo11n.pt")

try:
    while True:
        # 获取底部摄像头图像
        responses = client.simGetImages([
            airsim.ImageRequest("BottomCamera", airsim.ImageType.Scene, False, False)
        ])

        # 检查响应
        if responses and responses[0].width > 0 and responses[0].height > 0:
            # 将图像数据转换为 NumPy 数组
            img1d = np.frombuffer(responses[0].image_data_uint8, dtype=np.uint8)
            img_rgb = img1d.reshape(responses[0].height, responses[0].width, 3)
        #     img_bgr = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGR)  # 转换为 BGR 格式供 OpenCV 使用

            # 使用 YOLO 进行目标检测
            results = model(img_rgb)

            # 提取可视化结果图像
            annotated_frame = results[0].plot()  # 使用 YOLO 提供的 plot 方法

            # 显示图像
            cv2.imshow("Bottom Camera", annotated_frame)

            # 按下 'q' 键退出
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        else:
            print("Failed to retrieve image. Retrying...")
            time.sleep(0.1)  # 避免过多的无效请求

except KeyboardInterrupt:
    print("Stopped by user")

finally:
    # 释放资源
    cv2.destroyAllWindows()
    print("Exited.")
